package com.example.ai.agent;


import com.example.ai.agent.model.AgentState;
import io.micrometer.common.util.StringUtils;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

@Slf4j
@Data
public abstract class BaseAgent {
    //核心属性
    private String name;

    //提示信息
    private String systemPrompt;
    private String nextStepPrompt;

    //状态
    private AgentState state = AgentState.IDLE;

    //执行控制
    private int maxStep = 10;
    private int currentStep = 0;

    private ChatClient chatClient;

    //消息列表
    private List<Message> messagesList = new ArrayList<>();


    /**
     * 执行代理
     *
     * @param userPrompt 用户消息
     *                   return
     */
    public String run(String userPrompt) {
        if (this.state != AgentState.IDLE) {
            throw new RuntimeException("connot run agent from state: " + this.state);
        }
        if (StringUtils.isBlank(userPrompt)) {
            throw new RuntimeException("connot run agent with empty user Prompt");
        }
        //更改为运行状态
        this.state = AgentState.RUNNING;

        //上下文记录
        messagesList.add(new UserMessage(userPrompt));
        //结果保存
        List<String> results = new ArrayList<>();

        try { //执行循环
            for (int i = 0; i < maxStep && state != AgentState.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("ExecuTing Step {} of {}", stepNumber, maxStep);
                //执行
                String stepResult = step();
                //结果记录
                String result = "Step " + currentStep + ":" + stepResult;
            }
            if (currentStep >= maxStep) {
                state = AgentState.FINISHED;
                results.add("Terminated : Reached Max Step(" + maxStep + ")");
            }
            return String.join("\n", results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("Error runing agent " + e);
            return "Error runing agent: " + e.getMessage();
        } finally {
            //状态恢复
            this.cleanUp();
        }
    }

    public SseEmitter runSteam(String userPrompt){
        //5分钟时间等待
        SseEmitter sseEmitter = new SseEmitter(5 * 60 * 1000L);
        //使用线程异步处理 防止阻塞主流程
        CompletableFuture.runAsync(() -> {
            try {
                if (this.state != AgentState.IDLE) {
                    sseEmitter.send("Connot run agent from state: " + this.state);
                    sseEmitter.complete();
                    return;
                }
                if (StringUtils.isBlank(userPrompt)) {
                    sseEmitter.send("Connot run agent with empty user Prompt");
                    sseEmitter.complete();
                    return;
                }
            } catch (Exception e) {
               sseEmitter.completeWithError(e);
            }
            //执行代理逻辑更新状态
            this.state = AgentState.RUNNING;
            //记录上下文（用户消息）
            messagesList.add(new UserMessage(userPrompt));
            ArrayList<Object> results = new ArrayList<>();
            try {
                for (int i = 0; i < maxStep && state != AgentState.FINISHED; i++) {
                    int stepNumber = i + 1;
                    currentStep = stepNumber;
                    log.info("ExecuTing Step {} of {}", stepNumber, maxStep);
                    //执行
                    String stepResult = step();
                    //结果记录
                    String result = "Step " + currentStep + ":" + stepResult;
                    results.add(result);
                    //发送sse消息
                    sseEmitter.send(result);
                }
                if (currentStep >= maxStep) {
                    state = AgentState.FINISHED;
                    results.add("Terminated : Reached Max Steps(" + maxStep + ")");
                    sseEmitter.send("Terminated : Reached Max Steps(" + maxStep + ")");
                }
                sseEmitter.complete();
            }catch (Exception e) {
                state = AgentState.ERROR;
                log.error("Error runing agent " + e);
                try {
                    sseEmitter.send("Error running agent: " + e.getMessage());
                    sseEmitter.complete();
                }catch (Exception ex){
                    sseEmitter.completeWithError(ex);
                }
            }
            finally {
                this.cleanUp();
            }
        });
        //设置超时回报
        sseEmitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            this.cleanUp();
            log.warn("Timeout running agent");
        });
        //返回ssemitter
        sseEmitter.onCompletion(() -> {
         if (this.state == AgentState.RUNNING)  {
             this.state = AgentState.FINISHED;
         }
         this.cleanUp();
            log.info("Completed running agent");
        });
        return sseEmitter;
    }

    protected void cleanUp() {

    }

    /**
     * 定义单个逻辑
     *
     * @return
     */
    public abstract String step();


}
